package com.wolfking.spi;

import org.springframework.util.StringUtils;

import java.util.concurrent.Callable;

/**
 * spi bean的上下文和ThreadLocal
 *
 * @author 赵伟伟(wolfking)
 * created on 2019-09-29 13:33
 */
public class SpiApplicationContext {
    /**
     * spi的threadLocal,使用InheritableThreadLocal,复制父线程的ThreadLocal
     */
    private static ThreadLocal<String> spiCodeLocal = new InheritableThreadLocal<>();

    /**
     * 获取当前线程的 SpiCode
     *
     * @return
     */
    public static String getSpiCode() {
        return spiCodeLocal.get();
    }

    /**
     * 设置当前线程的SpiCode
     *
     * @param spiCode
     */
    public static void setSpiCode(String spiCode) {
        spiCodeLocal.set(spiCode);
    }

    /**
     * 清除当前线程的SpiCode
     *
     * @return
     */
    public static String cleanCode() {
        String code = getSpiCode();
        spiCodeLocal.remove();
        return code;
    }

    /**
     * 构建多线程的spiRunnable
     *
     * @param delegate 可执行的Runnable
     * @return
     */
    public static Runnable spiWrap(Runnable delegate) {
        /*先获取spiCode*/
        String spiCode = getSpiCode();
        return () -> {
            try {
                /*切换线程执行先设置spiCode*/
                if (!StringUtils.isEmpty(spiCode)) {
                    setSpiCode(spiCode);
                }
                delegate.run();
            } finally {
                /*执行完成之后清除spiCode，避免内存溢出*/
                cleanCode();
            }
        };
    }

    /**
     * 构建多线程的spiCallable
     *
     * @param delegate 原callable
     * @param <V>      返回参数的泛型规约
     * @return spi的Callable
     */
    public <V> Callable<V> spiWrap(Callable<V> delegate) {
        String spiCode = getSpiCode();
        return () -> {
            try {
                /*切换线程执行先设置spiCode*/
                if (!StringUtils.isEmpty(spiCode)) {
                    setSpiCode(spiCode);
                }
                return delegate.call();
            } finally {
                /*执行完成之后清除spiCode，避免内存溢出*/
                cleanCode();
            }
        };
    }
}
